{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ImageNet classification in a loop\n",
    "This notebook shows an example of ImageNet classification \n",
    "The network that is used for inference is a variant of DoReFaNet, whose topology is illustrated in the following picture.\n",
    "The pink layers are executed in the Programmable Logic at reduced precision (1 bit for weights, 2 bit for activations) while the other layers are executed in python.\n",
    "\n",
    "This notebook shows is a loop demo, please refer to [dorefanet-classification](./dorefanet-classification.ipynb) to run the classification on images choosen by the user\n",
    "\n",
    "![DoReFaNet topology](dorefanet-topology.svg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os, pickle, random\n",
    "from datetime import datetime\n",
    "from matplotlib import pyplot as plt\n",
    "from PIL import Image\n",
    "%matplotlib inline\n",
    "import IPython\n",
    "import numpy as np\n",
    "import cv2\n",
    "\n",
    "import qnn\n",
    "from qnn import Dorefanet\n",
    "from qnn import utils"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Instantiate a Classifier\n",
    "Creating a classifier will automatically download the bitstream onto the device, allocate memory buffers and load the network hyperparameters and weights. \n",
    "The neural network to be implemented is specified in a json file (*dorefanet-layers.json* in this example)\n",
    "The weights for the non-offloaded layers are also loaded in a numpy dictionary to be used for execution in python. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "classifier = Dorefanet()\n",
    "classifier.init_accelerator()\n",
    "net = classifier.load_network(json_layer=\"/opt/python3.6/lib/python3.6/site-packages/qnn/params/dorefanet-layers.json\")\n",
    "\n",
    "conv0_weights = np.load('/opt/python3.6/lib/python3.6/site-packages/qnn/params/dorefanet-conv0.npy', encoding=\"latin1\").item()\n",
    "fc_weights = np.load('/opt/python3.6/lib/python3.6/site-packages/qnn/params/dorefanet-fc-normalized.npy', encoding='latin1').item()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Get ImageNet Classes information\n",
    "Pick the random image from the imagenet-samples folder (image + correct class) and apply preprocessing transformation before inference. Thanks to the naming format adopted in these images (extracted from the validation set), the correct class is also displayed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open(\"imagenet-classes.pkl\", 'rb') as f:\n",
    "    classes = pickle.load(f)\n",
    "    names = dict((k, classes[k][1].split(',')[0]) for k in classes.keys())\n",
    "    synsets = dict((classes[k][0], classes[k][1].split(',')[0]) for k in classes.keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Launch demo\n",
    "The loop will automatically pick a random image from the imagenet-test folder (image + correct class) and perform the whole classification. Use the \"interrupt kernel\" button on top to stop the demo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "img_folder = './imagenet-samples/'\n",
    "\n",
    "conv0_W = conv0_weights['conv0/W']\n",
    "conv0_T = conv0_weights['conv0/T']\n",
    "# FC Layer 0\n",
    "fc0_W = fc_weights['fc0/Wn']\n",
    "fc0_b = fc_weights['fc0/bn']\n",
    "# FC Layer 1\n",
    "fc1_W = fc_weights['fc1/Wn']\n",
    "fc1_b = fc_weights['fc1/bn']\n",
    "# FC Layer 2\n",
    "fct_W = fc_weights['fct/W']\n",
    "\n",
    "in_dim = net['conv0']['output'][1]\n",
    "in_ch = net['conv0']['output'][0]\n",
    "out_dim = net['merge4']['output_dim']\n",
    "out_ch = net['merge4']['output_channels']\n",
    "\n",
    "conv_output = classifier.get_accel_buffer(out_ch, out_dim)\n",
    "\n",
    "while (1):\n",
    "    for image_name in os.listdir(img_folder):\n",
    "        img_file = os.path.join(img_folder, image_name)\n",
    "        img, img_class = classifier.load_image(img_file)\n",
    "        im = Image.open(img_file)\n",
    "        # 1st convolutional layer execution, having as input the image and the trained parameters (weights)\n",
    "        conv0 = utils.conv_layer(img, conv0_W, stride=4)\n",
    "        # The result in then quantized to 2 bits representation for the subsequent HW offload\n",
    "        conv0 = utils.threshold(conv0, conv0_T)\n",
    "        \n",
    "        conv_input = classifier.prepare_buffer(conv0)\n",
    "        \n",
    "        # Compute offloaded convolutional layers\n",
    "        classifier.inference(conv_input, conv_output)\n",
    "        fc_in = classifier.postprocess_buffer(conv_output)\n",
    "        \n",
    "        # Normalize results\n",
    "        fc_input = fc_in / np.max(fc_in)\n",
    "        fc0_out = utils.fully_connected(fc_input, fc0_W, fc0_b)\n",
    "        fc0_out = utils.qrelu(fc0_out)\n",
    "        fc0_out = utils.quantize(fc0_out, 2)\n",
    "        fc1_out = utils.fully_connected(fc0_out, fc1_W, fc1_b)\n",
    "        fc1_out = utils.qrelu(fc1_out)\n",
    "        fct_b = np.zeros((fct_W.shape[1], ))\n",
    "        fct_out = utils.fully_connected(fc1_out, fct_W, fct_b)\n",
    "        # Softmax\n",
    "        out = utils.softmax(fct_out)\n",
    "        # Top-5 results\n",
    "        topn =  utils.get_topn_indexes(out, 5)  \n",
    "        x_pos = np.arange(len(topn))\n",
    "        plt.barh(x_pos, out[topn], height=0.4, color='g', zorder=3)\n",
    "        plt.yticks(x_pos, [names[k] for k in topn])\n",
    "        plt.gca().invert_yaxis()\n",
    "        plt.xlim([0,1])\n",
    "        plt.grid(zorder=0)\n",
    "        IPython.display.clear_output(1)\n",
    "        display(im)\n",
    "        plt.show()\n",
    "        if img_class in synsets.keys():\n",
    "            print(\"Image class: {:>5}\\nPredictions:\".format(synsets[img_class]))\n",
    "            for k in topn: print(\"class:{0:>15}\\tprobability:{1:>8.2%}\".format(names[k].lower(), out[k]))\n",
    "            if synsets[img_class] in (names[k] for k in topn):\n",
    "                print(\"\\nMatch!\")\n",
    "        else:    \n",
    "            for k in topn: print(\"class:{0:>20}\\tprobability:{1:>8.2%}\".format(names[k].lower(), out[k]))\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reset the device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "classifier.deinit_accelerator()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from pynq import Xlnk\n",
    "\n",
    "xlnk = Xlnk();\n",
    "xlnk.xlnk_reset()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
